// SPDX-License-Identifier: GPL-2.0

#include <linux/wait.h>
#include <linux/compiler.h>
#include <linux/kernel.h>
#include <linux/slab.h>
#include <linux/err.h>
#include "dim_core_mem.h"

static inline unsigned long pow_of_2(unsigned int num)
{
	return 1 << num;
}

static inline bool is_pow_of_2(unsigned long num)
{
	if (num == 0)
		return false;

	return (num & (num - 1)) == 0;
}

static inline unsigned long get_max(unsigned long num1, unsigned long num2)
{
	return num1 > num2 ? num1 : num2;
}

static inline unsigned int right_child(unsigned int num)
{
	return (num << 1) + 1;
}

static inline unsigned int left_child(unsigned int num)
{
	return num << 1;
}

static inline unsigned int parent(unsigned int num)
{
	return num >> 1;
}

static unsigned int order_of_2(unsigned long num)
{
	unsigned int order = 0;

	while ((num = num >> 1) != 0)
		order++;

	return order;
}


static unsigned long up_to_pow_of_2(unsigned long num)
{
	unsigned int order = 1;

	if (is_pow_of_2(num))
		return num;

	while ((num = num >> 1) != 0)
		order++;

	return 1 << order;
}

/**
 * struct dim_buddy - buddy memory manage module implemented by binary tree
 * @size: all usable memory
 * @min_alloc_size: minimum alloc size of memory
 * @data: pre-alloc memory (several pages)
 * @node: binary-tree array, the value of node is the max usable memory size
 * @max_node_num: the length of the binary-tree array
 * @list: list struct for muti buddy modules
 */
struct dim_buddy {
	unsigned long size;
	unsigned long min_alloc_size;
	unsigned char *data;
	unsigned long *node;
	unsigned int max_node_num;
	struct list_head list;
};

#define DIM_BUDDY_MAX_GET_PAGES_ORDER 12
#define DIM_BUDDY_MIN_ALLOC_SIZE 32

/**
 * dim_buddy_create() - create a new buddy module
 * @order: the order of pages alloc
 * @min_alloc_size: minimum alloc size of memory
 *
 * Returns: a pointer to dim_buddy struct errno on failure (via ERR_PTR).
 */
static struct dim_buddy *dim_buddy_create(unsigned int order,
					  unsigned int min_alloc_size)
{
	unsigned int i;
	unsigned long size, node_size;
	struct dim_buddy *buddy = NULL;
	unsigned char *p = NULL;

	/* 1. check the request */
	if (order > DIM_BUDDY_MAX_GET_PAGES_ORDER)
		return ERR_PTR(-EINVAL);

	size = pow_of_2(order) * PAGE_SIZE;
	if (min_alloc_size > size)
		return ERR_PTR(-EINVAL);

	if (min_alloc_size < DIM_BUDDY_MIN_ALLOC_SIZE)
		min_alloc_size = DIM_BUDDY_MIN_ALLOC_SIZE;

	/* 2. pre-alloc pages */
	p = (unsigned char *)__get_free_pages(GFP_KERNEL, order);
	if (p == NULL)
		return ERR_PTR(-ENOMEM);

	/* 3. init dim_buddy struct */
	buddy = kmalloc(sizeof(struct dim_buddy), GFP_KERNEL);
	if (buddy == NULL) {
		free_pages((uintptr_t)p, order);
		return ERR_PTR(-ENOMEM);
	}

	buddy->data = p;
	buddy->size = size;
	buddy->min_alloc_size = up_to_pow_of_2(min_alloc_size);

	/* 4. init the buddy binary tree, the root index is 1 */
	buddy->max_node_num = ((buddy->size / buddy->min_alloc_size) << 1) - 1;
	buddy->node = kmalloc((buddy->max_node_num + 1) *
		sizeof(unsigned long), GFP_KERNEL);
	if (buddy->node == NULL) {
		free_pages((uintptr_t)p, order);
		kfree(buddy);
		return ERR_PTR(-ENOMEM);
	}

	node_size = buddy->size << 1;
	for (i = 1; i <= buddy->max_node_num; i++) {
		if (is_pow_of_2(i))
			node_size = (node_size >> 1);

		buddy->node[i] = node_size;
	}

	return buddy;
}

static void dim_buddy_destroy(struct dim_buddy *buddy)
{
	if (buddy == NULL)
		return;

	if (buddy->data != NULL)
		free_pages((unsigned long)buddy->data,
			order_of_2(buddy->size / PAGE_SIZE));

	if (buddy->node != NULL)
		kfree(buddy->node);

	kfree(buddy);
}

static void *dim_buddy_alloc(struct dim_buddy *buddy, unsigned long size)
{
	unsigned int cur = 1;
	unsigned int i;
	unsigned long f_size, node_size;

	if (size < buddy->min_alloc_size)
		size = buddy->min_alloc_size;

	f_size = up_to_pow_of_2(size);
	if (buddy->node[1] < f_size)
		return NULL; /* no memory */

	/* find suitable node */
	while (cur <= buddy->max_node_num) {
		if (buddy->node[cur] >= f_size) {
			/* find ok */
			if (right_child(cur) > buddy->max_node_num)
				break; /* last layer */

			if (buddy->node[left_child(cur)] >= f_size) {
				cur = left_child(cur);
				continue;
			}

			if (buddy->node[right_child(cur)] >= f_size) {
				cur = right_child(cur);
				continue;
			}

			break;
		}
	}

	if (cur > buddy->max_node_num)
		return NULL; /* shouldn't run here */

	node_size = buddy->node[cur];
	buddy->node[cur] = 0; /* mark the node */

	/* trace back */
	for (i = parent(cur); i >= 1; i = parent(i))
		buddy->node[i] = get_max(buddy->node[left_child(i)],
			buddy->node[right_child(i)]);

	/* calculate and return offset */
	return &buddy->data[cur * node_size - buddy->size];
}

/*
	note: data must be allocated by this buddy module
*/
static void dim_buddy_free(struct dim_buddy *buddy, void *data)
{
	long offset;
	int cur, i;
	unsigned long node_size;

	if (buddy == NULL || data == NULL)
		return;

	offset = (uintptr_t)data - (uintptr_t)buddy->data;
	if (offset < 0 || offset >= buddy->size)
		return;

	node_size = buddy->min_alloc_size; /* find from min size */
	cur = (buddy->size + offset) / node_size;

	for (; cur >= 1; cur = parent(cur)) {
		if (buddy->node[cur] == 0) {
			buddy->node[cur] = node_size; /* recorver the node */
			break;
		}

		node_size = (node_size << 1);
	}

	/* trace back */
	for (i = parent(cur); i >= 1; i = parent(i)) {
		node_size = (node_size << 1);
		if (node_size == (buddy->node[left_child(i)] + buddy->node[right_child(i)]))
			buddy->node[i] = node_size; /* merge memory */
		else
			buddy->node[i] = get_max(buddy->node[left_child(i)],
						buddy->node[right_child(i)]);
	}
}

static struct list_head buddy_list;
static int buddy_num;

int dim_mem_init(void)
{
	int i, ret;
	struct dim_buddy *buddy = NULL;
	struct list_head *node = NULL;

	INIT_LIST_HEAD(&buddy_list);
	for (i = 0; i < DIM_INIT_BUDDY_NUMBER; i++) {
		buddy = dim_buddy_create(DIM_GET_PAGES_ORDER,
			DIM_MIN_ALLOC_SIZE);
		if (IS_ERR(buddy)) {
			ret = PTR_ERR(buddy);
			goto err;
		}

		list_add_tail(&buddy->list, &buddy_list);
	}

	buddy_num = DIM_INIT_BUDDY_NUMBER;
	return 0;

err:
	list_for_each(node, &buddy_list) {
		buddy = list_entry(node, struct dim_buddy, list);
		dim_buddy_destroy(buddy);
	}

	return ret;
}

void dim_mem_destroy(void)
{
	struct dim_buddy *buddy = NULL;
	struct list_head *node = NULL;
	struct list_head *next = NULL;

	list_for_each_safe(node, next, &buddy_list) {
		buddy = list_entry(node, struct dim_buddy, list);
		dim_buddy_destroy(buddy);
	}
}

void *dim_mem_alloc(unsigned long size)
{
	struct dim_buddy *buddy = NULL;
	struct list_head *node = NULL;
	void *data = NULL;

	if (size > DIM_MAX_ALLOC_SIZE)
		return NULL;

	list_for_each(node, &buddy_list) {
		buddy = list_entry(node, struct dim_buddy, list);
		if (buddy->node[1] >= size) {
			data = dim_buddy_alloc(buddy, size);
			if (data == NULL)
				continue; /* shouldn't run here */
			return data;
		}
	}

	/* no memory */
	if (buddy_num >= DIM_MAX_BUDDY_NUMBER)
		return NULL;

	buddy = dim_buddy_create(DIM_GET_PAGES_ORDER, DIM_MIN_ALLOC_SIZE);
	if (buddy == NULL)
		return NULL;

	list_add_tail(&buddy->list, &buddy_list);
	buddy_num++;

	return dim_buddy_alloc(buddy, size);
}

void dim_mem_free(void *data)
{
	struct dim_buddy *buddy = NULL;
	struct list_head *node = NULL;

	if (data == NULL)
		return;

	list_for_each(node, &buddy_list) {
		buddy = list_entry(node, struct dim_buddy, list);
		if ((uintptr_t)data >= (uintptr_t)buddy->data &&
		    (uintptr_t)data - (uintptr_t)buddy->data < buddy->size) {
			dim_buddy_free(buddy, data);
			return;
		}
	}
}

int dim_mem_block_traverse(dim_mem_iterator f, void *data)
{
	struct dim_buddy *buddy = NULL;
	struct list_head *node = NULL;
	int ret;

	if (f == NULL)
		return -EINVAL;

	list_for_each(node, &buddy_list) {
		buddy = list_entry(node, struct dim_buddy, list);
		ret = f(buddy->data, buddy->size, data);
		if (ret != 0)
			return ret;
	}

	return 0;
}

unsigned long dim_mem_get_block_counts(void)
{
	unsigned long counts = 0;
	struct list_head *node = NULL;

	list_for_each(node, &buddy_list)
		counts++;

	return counts;
}
